"""
MessagePack 序列化器实现
提供紧凑的二进制序列化，无需schema定义
"""

import time
import uuid
import zlib
import lz4.frame
from typing import Dict, Any, Optional, Union, Callable
from dataclasses import dataclass, asdict
from datetime import datetime, date
from decimal import Decimal

try:
    import msgpack

    MSGPACK_AVAILABLE = True
except ImportError:
    MSGPACK_AVAILABLE = False
    print("警告: msgpack 未安装，请运行: pip install msgpack")


@dataclass
class MessagePackStats:
    """MessagePack序列化统计信息"""
    total_packed: int = 0
    total_unpacked: int = 0
    total_bytes: int = 0
    compressed_bytes: int = 0
    avg_pack_time: float = 0.0
    avg_unpack_time: float = 0.0
    compression_ratio: float = 0.0
    error_count: int = 0


class CustomTypeRegistry:
    """自定义类型注册器"""

    def __init__(self):
        self.encoders = {}
        self.decoders = {}
        self.type_codes = {}
        self.next_code = 1

    def register_type(self, type_class, encoder: Callable, decoder: Callable, type_code: int = None):
        """注册自定义类型"""
        if type_code is None:
            type_code = self.next_code
            self.next_code += 1

        self.encoders[type_class] = (type_code, encoder)
        self.decoders[type_code] = decoder
        self.type_codes[type_class] = type_code

    def encode_custom(self, obj):
        """编码自定义类型"""
        obj_type = type(obj)
        if obj_type in self.encoders:
            type_code, encoder = self.encoders[obj_type]
            encoded_data = encoder(obj)
            return msgpack.ExtType(type_code, encoded_data)
        raise TypeError(f"无法序列化类型: {obj_type}")

    def decode_custom(self, code, data):
        """解码自定义类型"""
        if code in self.decoders:
            decoder = self.decoders[code]
            return decoder(data)
        raise ValueError(f"未知的类型代码: {code}")


class MessagePackSerializer:
    """MessagePack 序列化器"""

    def __init__(self, use_compression: bool = True, compression_algorithm: str = 'lz4',
                 compression_threshold: int = 1024, strict_types: bool = False):
        if not MSGPACK_AVAILABLE:
            raise ImportError("msgpack 未安装，请运行: pip install msgpack")

        self.use_compression = use_compression
        self.compression_algorithm = compression_algorithm
        self.compression_threshold = compression_threshold
        self.strict_types = strict_types
        self.stats = MessagePackStats()

        # 自定义类型注册器
        self.type_registry = CustomTypeRegistry()
        self._register_default_types()

        # MessagePack配置
        self.pack_options = {
            'use_bin_type': True,
            'strict_types': strict_types,
            'default': self.type_registry.encode_custom,
            'autoreset': True,
            'use_single_float': False,
        }

        self.unpack_options = {
            'raw': False,
            'strict_map_key': False,
            'ext_hook': self.type_registry.decode_custom,
            #'max_buffer_size': 100 * 1024 * 1024,  # 100MB
        }

    def _register_default_types(self):
        """注册默认的自定义类型"""
        # 日期时间类型
        self.type_registry.register_type(
            datetime,
            lambda dt: dt.isoformat().encode('utf-8'),
            lambda data: datetime.fromisoformat(data.decode('utf-8')),
            1
        )

        self.type_registry.register_type(
            date,
            lambda d: d.isoformat().encode('utf-8'),
            lambda data: date.fromisoformat(data.decode('utf-8')),
            2
        )

        # Decimal类型
        self.type_registry.register_type(
            Decimal,
            lambda d: str(d).encode('utf-8'),
            lambda data: Decimal(data.decode('utf-8')),
            3
        )

        # bytes类型（如果需要特殊处理）
        self.type_registry.register_type(
            bytes,
            lambda b: b,
            lambda data: data,
            4
        )

    def _compress_data(self, data: bytes) -> bytes:
        """压缩数据"""
        if not self.use_compression or len(data) < self.compression_threshold:
            return data

        if self.compression_algorithm == 'lz4':
            return lz4.frame.compress(data)
        elif self.compression_algorithm == 'zlib':
            return zlib.compress(data, level=6)
        else:
            return data

    def _decompress_data(self, data: bytes) -> bytes:
        """解压缩数据"""
        if not self.use_compression:
            return data

        try:
            if self.compression_algorithm == 'lz4':
                return lz4.frame.decompress(data)
            elif self.compression_algorithm == 'zlib':
                return zlib.decompress(data)
        except Exception:
            # 如果解压缩失败，可能是未压缩的数据
            return data

        return data

    def pack(self, obj: Any) -> bytes:
        """打包对象为MessagePack格式"""
        start_time = time.perf_counter()

        try:
            # 预处理对象
            processed_obj = self._preprocess_object(obj)

            # MessagePack序列化
            packed_data = msgpack.packb(processed_obj, **self.pack_options)

            # 可选压缩
            original_size = len(packed_data)
            compressed_data = self._compress_data(packed_data)

            # 添加压缩标记
            if len(compressed_data) < original_size:
                # 压缩有效，添加压缩标记
                result = b'\x01' + compressed_data  # 0x01表示已压缩
                self.stats.compressed_bytes += len(result)
            else:
                # 压缩无效，使用原始数据
                result = b'\x00' + packed_data  # 0x00表示未压缩

            # 更新统计
            self.stats.total_packed += 1
            self.stats.total_bytes += len(result)

            return result

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_pack_time(duration)

    def unpack(self, data: bytes) -> Any:
        """从MessagePack格式解包对象"""
        start_time = time.perf_counter()

        try:
            if len(data) < 1:
                raise ValueError("数据太短")

            # 检查压缩标记
            compression_flag = data[0]
            payload = data[1:]

            # 解压缩（如果需要）
            if compression_flag == 0x01:
                payload = self._decompress_data(payload)

            # MessagePack反序列化
            unpacked_obj = msgpack.unpackb(payload, **self.unpack_options)

            # 后处理对象
            result = self._postprocess_object(unpacked_obj)

            # 更新统计
            self.stats.total_unpacked += 1

            return result

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_unpack_time(duration)

    def pack_stream(self, objects: list) -> bytes:
        """流式打包多个对象"""
        start_time = time.perf_counter()

        try:
            # 使用Packer进行流式打包
            packer = msgpack.Packer(**self.pack_options)

            packed_data = b''
            for obj in objects:
                processed_obj = self._preprocess_object(obj)
                packed_data += packer.pack(processed_obj)

            # 压缩整个流
            compressed_data = self._compress_data(packed_data)

            if len(compressed_data) < len(packed_data):
                result = b'\x01' + compressed_data
                self.stats.compressed_bytes += len(result)
            else:
                result = b'\x00' + packed_data

            # 更新统计
            self.stats.total_packed += len(objects)
            self.stats.total_bytes += len(result)

            return result

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_pack_time(duration)

    def unpack_stream(self, data: bytes) -> list:
        """流式解包多个对象"""
        start_time = time.perf_counter()

        try:
            if len(data) < 1:
                raise ValueError("数据太短")

            # 检查压缩标记
            compression_flag = data[0]
            payload = data[1:]

            # 解压缩（如果需要）
            if compression_flag == 0x01:
                payload = self._decompress_data(payload)

            # 使用Unpacker进行流式解包
            unpacker = msgpack.Unpacker(raw=False, **self.unpack_options)
            unpacker.feed(payload)

            objects = []
            for unpacked_obj in unpacker:
                result = self._postprocess_object(unpacked_obj)
                objects.append(result)

            # 更新统计
            self.stats.total_unpacked += len(objects)

            return objects

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_unpack_time(duration)

    def _preprocess_object(self, obj: Any) -> Any:
        """预处理对象（序列化前）"""
        if isinstance(obj, dict):
            # 处理字典
            return {k: self._preprocess_object(v) for k, v in obj.items()}
        elif isinstance(obj, (list, tuple)):
            # 处理列表和元组
            return [self._preprocess_object(item) for item in obj]
        elif hasattr(obj, '__dict__'):
            # 处理自定义对象
            return {
                '__class__': obj.__class__.__name__,
                '__module__': obj.__class__.__module__,
                '__data__': self._preprocess_object(obj.__dict__)
            }
        else:
            return obj

    def _postprocess_object(self, obj: Any) -> Any:
        """后处理对象（反序列化后）"""
        if isinstance(obj, dict):
            # 检查是否是自定义对象
            if '__class__' in obj and '__module__' in obj and '__data__' in obj:
                # 尝试重建自定义对象
                try:
                    module = __import__(obj['__module__'], fromlist=[obj['__class__']])
                    cls = getattr(module, obj['__class__'])
                    instance = cls.__new__(cls)
                    instance.__dict__.update(self._postprocess_object(obj['__data__']))
                    return instance
                except (ImportError, AttributeError):
                    # 如果无法重建，返回原始字典
                    return obj
            else:
                # 普通字典
                return {k: self._postprocess_object(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [self._postprocess_object(item) for item in obj]
        else:
            return obj

    def _update_pack_time(self, duration: float):
        """更新打包时间统计"""
        if self.stats.total_packed == 1:
            self.stats.avg_pack_time = duration
        else:
            alpha = 0.1
            self.stats.avg_pack_time = (
                    alpha * duration + (1 - alpha) * self.stats.avg_pack_time
            )

    def _update_unpack_time(self, duration: float):
        """更新解包时间统计"""
        if self.stats.total_unpacked == 1:
            self.stats.avg_unpack_time = duration
        else:
            alpha = 0.1
            self.stats.avg_unpack_time = (
                    alpha * duration + (1 - alpha) * self.stats.avg_unpack_time
            )

    def get_stats(self) -> Dict[str, Any]:
        """获取序列化统计信息"""
        stats_dict = asdict(self.stats)

        # 计算压缩比
        if self.stats.total_bytes > 0:
            stats_dict['compression_ratio'] = (
                    1.0 - (self.stats.compressed_bytes / self.stats.total_bytes)
            ) if self.stats.compressed_bytes > 0 else 0.0

        # 计算性能指标
        if self.stats.avg_pack_time > 0:
            stats_dict['pack_rate'] = 1.0 / self.stats.avg_pack_time
        if self.stats.avg_unpack_time > 0:
            stats_dict['unpack_rate'] = 1.0 / self.stats.avg_unpack_time

        # 添加配置信息
        stats_dict['config'] = {
            'use_compression': self.use_compression,
            'compression_algorithm': self.compression_algorithm,
            'compression_threshold': self.compression_threshold,
            'strict_types': self.strict_types,
        }

        return stats_dict

    def reset_stats(self):
        """重置统计信息"""
        self.stats = MessagePackStats()

    def benchmark(self, test_data: list, iterations: int = 1000) -> Dict[str, Any]:
        """性能基准测试"""
        import random

        # 重置统计
        self.reset_stats()

        # 准备测试数据
        test_objects = []
        for _ in range(iterations):
            data_item = random.choice(test_data)
            test_objects.append(data_item)

        # 打包基准测试
        start_time = time.perf_counter()
        packed_messages = []

        for obj in test_objects:
            packed = self.pack(obj)
            packed_messages.append(packed)

        pack_time = time.perf_counter() - start_time

        # 解包基准测试
        start_time = time.perf_counter()

        for packed in packed_messages:
            self.unpack(packed)

        unpack_time = time.perf_counter() - start_time

        # 计算结果
        total_size = sum(len(msg) for msg in packed_messages)

        return {
            'iterations': iterations,
            'pack_time': pack_time,
            'unpack_time': unpack_time,
            'pack_rate': iterations / pack_time,
            'unpack_rate': iterations / unpack_time,
            'total_size': total_size,
            'avg_message_size': total_size / iterations,
            'throughput_mbps': (total_size / (1024 * 1024)) / (pack_time + unpack_time),
            'stats': self.get_stats()
        }

    def register_custom_type(self, type_class, encoder: Callable, decoder: Callable):
        """注册自定义类型"""
        self.type_registry.register_type(type_class, encoder, decoder)

    def create_message(self, msg_type: str, data: Any, **metadata) -> Dict[str, Any]:
        """创建标准消息格式"""
        return {
            'header': {
                'type': msg_type,
                'timestamp': time.time(),
                'seq_id': int(time.time() * 1000000),
                'correlation_id': str(uuid.uuid4()),
                'metadata': metadata
            },
            'payload': data
        }


# 使用示例
if __name__ == "__main__":
    if MSGPACK_AVAILABLE:
        # 创建序列化器
        serializer = MessagePackSerializer(use_compression=True)

        # 测试基础数据类型
        test_data = {
            'string': 'Hello, MessagePack!',
            'number': 42,
            'float': 3.14159,
            'boolean': True,
            'null': None,
            'list': [1, 2, 3, 4, 5],
            'dict': {'nested': 'value'},
            'bytes': b'binary data',
            'datetime': datetime.now(),
        }

        # 打包
        packed = serializer.pack(test_data)
        print(f"打包数据: {len(packed)} bytes")

        # 解包
        unpacked = serializer.unpack(packed)
        print(f"解包结果: {unpacked}")

        # 测试流式处理
        test_objects = [
            {'id': i, 'data': f'message_{i}'}
            for i in range(10)
        ]

        stream_packed = serializer.pack_stream(test_objects)
        print(f"流式打包: {len(stream_packed)} bytes")

        stream_unpacked = serializer.unpack_stream(stream_packed)
        print(f"流式解包: {len(stream_unpacked)} objects")

        # 获取统计信息
        stats = serializer.get_stats()
        print(f"统计信息: {stats}")


        # 测试自定义类型
        @dataclass
        class CustomMessage:
            id: int
            content: str
            timestamp: datetime


        # 注册自定义类型
        def encode_custom_message(msg):
            return msgpack.packb({
                'id': msg.id,
                'content': msg.content,
                'timestamp': msg.timestamp.isoformat()
            })


        def decode_custom_message(data):
            unpacked = msgpack.unpackb(data, raw=False)
            return CustomMessage(
                id=unpacked['id'],
                content=unpacked['content'],
                timestamp=datetime.fromisoformat(unpacked['timestamp'])
            )


        serializer.register_custom_type(CustomMessage, encode_custom_message, decode_custom_message)

        # 测试自定义类型
        custom_msg = CustomMessage(1, "Hello Custom!", datetime.now())
        custom_packed = serializer.pack(custom_msg)
        custom_unpacked = serializer.unpack(custom_packed)

        print(f"自定义类型: {custom_unpacked}")

    else:
        print("请安装 msgpack: pip install msgpack")